import argparse
import logging
import os
import sys
from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import optim
from torch.utils.data import DataLoader, random_split
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm

from dice_loss import dice_coeff
from unet import UNet
from utils.dataset import BasicDataset

dir_img = '/data/xinlin/Segmentation/Carvana/train/'
dir_mask = '/data/xinlin/Segmentation/Carvana/train_masks/'
dir_checkpoint = 'checkpoints/'


def get_args():
    parser = argparse.ArgumentParser(description='Train the UNet on images and target masks')
    parser.add_argument('-e', '--epochs', metavar='E', type=int, default=5,
                        help='Number of epochs', dest='epochs')
    parser.add_argument('-b', '--batch-size', metavar='B', type=int, nargs='?', default=1,
                        help='Batch size', dest='batch_size')
    parser.add_argument('-l', '--learning-rate', metavar='LR', type=float, nargs='?', default=0.0001,
                        help='Learning rate', dest='lr')
    parser.add_argument('-f', '--load', dest='load', type=str, default=False,
                        help='Load model from a .pth file')
    parser.add_argument('-s', '--scale', dest='scale', type=float, default=0.5,
                        help='Downscaling factor of the images')
    parser.add_argument('-v', '--validation', dest='val', type=float, default=10.0,
                        help='Percent of the data that is used as validation (0-100)')
    return parser.parse_args()


class Trainer:

    def __init__(self, args, device, save_cp=True):
        self.args = args

        self.writer = SummaryWriter(comment=f'LR_{args.lr}_BS_{args.batch_size}_SCALE_{args.scale}')

        # Change here to adapt to your data
        # n_channels=3 for RGB images
        # n_classes is the number of probabilities you want to get per pixel
        #   - For 1 class and background, use n_classes=1
        #   - For 2 classes, use n_classes=1
        #   - For N > 2 classes, use n_classes=N
        self.net = UNet(n_channels=3, n_classes=1, bilinear=True)
        logging.info(f'Network:\n'
                     f'\t{self.net.n_channels} input channels\n'
                     f'\t{self.net.n_classes} output channels (classes)\n'
                     f'\t{"Bilinear" if self.net.bilinear else "Transposed conv"} upscaling')
        if args.load:
            self.net.load_state_dict(torch.load(args.load, map_location=device))
            logging.info(f'Model loaded from {args.load}')
        self.net.to(device=device)
        # faster convolutions, but more memory
        # cudnn.benchmark = True

        self.optimizer = optim.RMSprop(self.net.parameters(), lr=args.lr, weight_decay=1e-8, momentum=0.9)
        self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(self.optimizer,
                                                              'min' if self.net.n_classes > 1 else 'max',
                                                              patience=2)
        if self.net.n_classes > 1:
            self.criterion = nn.CrossEntropyLoss()
        else:
            self.criterion = nn.BCEWithLogitsLoss()

        dataset = BasicDataset(dir_img, dir_mask, args.scale, mask_suffix='_mask')
        n_val = int(len(dataset) * (args.val / 100))
        n_train = len(dataset) - n_val
        train, val = random_split(dataset, [n_train, n_val], torch.default_generator)
        self.train_loader = DataLoader(train, args.batch_size, shuffle=True, num_workers=1, pin_memory=True)
        self.val_loader = DataLoader(val, args.batch_size, shuffle=False, num_workers=1, pin_memory=True,
                                     drop_last=True)

        # train_net(self.net=net, epochs=args.epochs, batch_size=args.batchsize, lr=args.lr, device=device,
        #                   img_scale=args.scale, val_percent=args.val / 100)
        logging.info(f'''Starting training:
                        Epochs:          {args.epochs}
                        Batch size:      {args.batchsize}
                        Learning rate:   {args.lr}
                        Training size:   {n_train}
                        Validation size: {n_val}
                        Checkpoints:     {save_cp}
                        Device:          {device.type}
                        Images scaling:  {args.scale}
                    ''')
        self.global_step = 0

    def train(self, epoch, device):
        self.net.train()

        epoch_loss = 0
        tbar = tqdm(self.train_loader)
        n_train = len(self.train_loader)
        # with tqdm(total=n_train, desc=f'Epoch {epoch + 1}/{self.args.epochs}', unit='img') as pbar:
        for i, batch in enumerate(tbar):
            imgs = batch['image']
            true_masks = batch['mask']
            assert imgs.shape[1] == self.net.n_channels, \
                f'Network has been defined with {self.net.n_channels} input channels, ' \
                f'but loaded images have {imgs.shape[1]} channels. Please check that ' \
                'the images are loaded correctly.'

            imgs = imgs.to(device=device, dtype=torch.float32)
            mask_type = torch.float32 if self.net.n_classes == 1 else torch.long
            true_masks = true_masks.to(device=device, dtype=mask_type)

            masks_pred = self.net(imgs)
            loss = self.criterion(masks_pred, true_masks)
            epoch_loss += loss.item()
            self.writer.add_scalar('Loss/train', loss.item(), self.global_step)

            # pbar.set_postfix(**{'loss (batch)': loss.item()})

            self.optimizer.zero_grad()
            loss.backward()
            nn.utils.clip_grad_value_(self.net.parameters(), 0.1)
            self.optimizer.step()

            # pbar.update(imgs.shape[0])
            desc = 'Train: [Epoch: {}/{}, LR: {:.5f}], loss (batch)={:.4f}'.format(epoch + 1, self.args.epochs,
                                                                                   self.args.lr, loss.item())
            tbar.set_description(desc)
            self.global_step += 1

            self.update_tensorboard(imgs, true_masks, masks_pred)

    def update_tensorboard(self, imgs, true_masks, masks_pred):
        n_train = len(self.train_loader)
        if self.global_step % (n_train // (10 * self.args.batch_size)) == 0:
            for tag, value in self.net.named_parameters():
                tag = tag.replace('.', '/')
                self.writer.add_histogram('weights/' + tag, value.data.cpu().numpy(), self.global_step)
                self.writer.add_histogram('grads/' + tag, value.grad.data.cpu().numpy(), self.global_step)
            self.writer.add_scalar('learning_rate', self.optimizer.param_groups[0]['lr'], self.global_step)

            self.writer.add_images('images', imgs, self.global_step)
            if self.net.n_classes == 1:
                self.writer.add_images('masks/true', true_masks, self.global_step)
                self.writer.add_images('masks/pred', torch.sigmoid(masks_pred) > 0.5, self.global_step)

    def eval(self, device):
        """Evaluation without the densecrf with the dice coefficient"""
        self.net.eval()

        mask_type = torch.float32 if self.net.n_classes == 1 else torch.long
        n_val = len(self.val_loader)  # the number of batch
        n_train = len(self.train_loader)
        tot = 0
        with tqdm(total=n_val, desc='Validation round', unit='batch', leave=False) as pbar:
            for batch in self.val_loader:
                imgs, true_masks = batch['image'], batch['mask']
                imgs = imgs.to(device=device, dtype=torch.float32)
                true_masks = true_masks.to(device=device, dtype=mask_type)

                with torch.no_grad():
                    mask_pred = self.net(imgs)

                if self.net.n_classes > 1:
                    tot += F.cross_entropy(mask_pred, true_masks).item()
                else:
                    pred = torch.sigmoid(mask_pred)
                    pred = (pred > 0.5).float()
                    tot += dice_coeff(pred, true_masks).item()
                pbar.update()

        val_score = tot / n_val
        self.scheduler.step(val_score)
        if self.net.n_classes > 1:
            logging.info('Validation cross entropy: {}'.format(val_score))
            self.writer.add_scalar('Loss/test', val_score, self.global_step)
        else:
            logging.info('Validation Dice Coeff: {}'.format(val_score))
            self.writer.add_scalar('Dice/test', val_score, self.global_step)

    def close(self, epoch):
        Path(dir_checkpoint).mkdir(exist_ok=True)
        logging.info('Created checkpoint directory')
        torch.save(self.net.state_dict(), dir_checkpoint + f'CP_epoch{epoch + 1}.pth')
        logging.info(f'Checkpoint {epoch + 1} saved !')
        self.writer.close()


def main():
    logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')
    args = get_args()
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    logging.info(f'Using device {device}')
    trainer = Trainer(args, device)
    try:
        epoch = 0
        for epoch in range(args.epochs):
            trainer.train(epoch, device)
            trainer.eval(device)
        trainer.close(epoch)
    except KeyboardInterrupt:
        torch.save(trainer.net.state_dict(), 'INTERRUPTED.pth')
        logging.info('Saved interrupt')
        try:
            sys.exit(0)
        except SystemExit:
            os._exit(0)


if __name__ == '__main__':
    main()
